{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# HybridNonlinearFactorGraph"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/hybrid/doc/HybridNonlinearFactorGraph.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`HybridNonlinearFactorGraph` is a `HybridFactorGraph` specifically for representing **nonlinear** hybrid optimization problems. It can contain:\n",
        "\n",
        "*   `gtsam.NonlinearFactor` (including derived types like `PriorFactor`, `BetweenFactor`)\n",
        "*   `gtsam.DiscreteFactor` (including `DecisionTreeFactor`, `TableFactor`)\n",
        "*   `gtsam.HybridNonlinearFactor`\n",
        "\n",
        "This is the graph type you would typically build to model a system with both continuous states (potentially on manifolds) and discrete modes or decisions, where the relationships are nonlinear.\n",
        "\n",
        "The primary operation performed on a `HybridNonlinearFactorGraph` is `linearize`, which converts it into a `HybridGaussianFactorGraph` suitable for inference using methods like `eliminateSequential` or `eliminateMultifrontal`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "\n",
        "from gtsam import (\n",
        "    HybridNonlinearFactorGraph, HybridNonlinearFactor,\n",
        "    PriorFactorPose2, BetweenFactorPose2, Pose2, Point3,\n",
        "    DecisionTreeFactor, Values, DiscreteValues, HybridValues,\n",
        "    HybridGaussianFactorGraph\n",
        ")\n",
        "from gtsam.symbol_shorthand import X, D"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialization and Adding Factors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HybridNonlinearFactorGraph Contents:\n",
            "HybridNonlinearFactorGraph\n",
            " \n",
            "size: 3\n",
            "factor 0: PriorFactor on x0\n",
            "  prior mean:  (0, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.05];\n",
            "\n",
            "factor 1:  f[ (d0,2), ]\n",
            " Choice(d0) \n",
            " 0 Leaf  0.8\n",
            " 1 Leaf  0.2\n",
            "\n",
            "factor 2:  Hybrid [x0 x1; d0]\n",
            "HybridNonlinearFactor\n",
            " Choice(d0) \n",
            " 0 Leaf (0) BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.0174532925];\n",
            "\n",
            " 1 Leaf (0) BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.5; 0.5; 0.174532925];\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "hnfg = HybridNonlinearFactorGraph()\n",
        "\n",
        "# 1. Add a Nonlinear Factor (Prior)\n",
        "prior_noise = gtsam.noiseModel.Diagonal.Sigmas(Point3(0.1, 0.1, 0.05))\n",
        "hnfg.add(PriorFactorPose2(X(0), Pose2(0, 0, 0), prior_noise))\n",
        "\n",
        "# 2. Add a Discrete Factor (Prior on mode)\n",
        "dk0 = (D(0), 2) # Binary mode\n",
        "hnfg.add(DecisionTreeFactor([dk0], \"0.8 0.2\")) # P(D0=0)=0.8\n",
        "\n",
        "# 3. Add a HybridNonlinearFactor (Mode-dependent odometry)\n",
        "# Mode 0 (Grippy): Smaller noise\n",
        "noise0 = gtsam.noiseModel.Diagonal.Sigmas(Point3(0.1, 0.1, np.radians(1)))\n",
        "odom0 = BetweenFactorPose2(X(0), X(1), Pose2(1.0, 0, 0), noise0)\n",
        "# Mode 1 (Slippery): Larger noise\n",
        "noise1 = gtsam.noiseModel.Diagonal.Sigmas(Point3(0.5, 0.5, np.radians(10)))\n",
        "odom1 = BetweenFactorPose2(X(0), X(1), Pose2(1.0, 0, 0), noise1)\n",
        "# Assume equal probability for modes within this factor's context\n",
        "hybrid_odom = HybridNonlinearFactor(dk0, [odom0, odom1])\n",
        "hnfg.add(hybrid_odom)\n",
        "\n",
        "print(\"HybridNonlinearFactorGraph Contents:\")\n",
        "hnfg.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Error Calculation (`error`, `errorTree`, `discretePosterior`)\n",
        "\n",
        "Evaluates the combined error (nonlinear factors) and negative log probability (discrete factors) of the graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluating error with HybridValues:\n",
            "HybridValues: \n",
            "  Continuous: 0 elements\n",
            "  Discrete: (d0, 0)\n",
            "  Nonlinear\n",
            "Values with 2 values:\n",
            "Value x0: (gtsam::Pose2)\n",
            "(0.05, 0.05, 0.0174532925)\n",
            "\n",
            "Value x1: (gtsam::Pose2)\n",
            "(1, -0.05, 0)\n",
            "\n",
            "\n",
            "Total graph error: 1.8480600597872188\n"
          ]
        }
      ],
      "source": [
        "# Requires HybridValues for evaluation\n",
        "hybrid_values = HybridValues()\n",
        "\n",
        "# Continuous/Nonlinear Values\n",
        "nl_vals = Values()\n",
        "nl_vals.insert(X(0), Pose2(0.05, 0.05, np.radians(1))) # Near prior mean\n",
        "nl_vals.insert(X(1), Pose2(1.0, -0.05, np.radians(0))) # Near odom mean\n",
        "hybrid_values.insert(nl_vals)\n",
        "\n",
        "# Discrete Values\n",
        "disc_vals = DiscreteValues()\n",
        "disc_vals[D(0)] = 0 # Select Grippy mode (lower error)\n",
        "hybrid_values.insert(disc_vals)\n",
        "\n",
        "print(\"\\nEvaluating error with HybridValues:\")\n",
        "hybrid_values.print()\n",
        "\n",
        "# --- error(HybridValues) ---\n",
        "total_error = hnfg.error(hybrid_values)\n",
        "print(f\"\\nTotal graph error: {total_error}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "65ca9727",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Error Tree (Nonlinear errors across modes):\n",
            "ErrorTree\n",
            " Choice(d0) \n",
            "ErrorTree\n",
            " 0 Leaf 1.8480601\n",
            "ErrorTree\n",
            " 1 Leaf 1.9579211\n",
            "\n",
            "Discrete Posterior P(M | X=x):\n",
            "ErrorTree\n",
            " Choice(d0) \n",
            "ErrorTree\n",
            " 0 Leaf 0.52743767\n",
            "ErrorTree\n",
            " 1 Leaf 0.47256233\n"
          ]
        }
      ],
      "source": [
        "# --- errorTree(Values) ---\n",
        "# Calculates error for Nonlinear & HybridNonlinear factors across all discrete modes\n",
        "# Ignores purely DiscreteFactors for this calculation.\n",
        "cont_values = hybrid_values.nonlinear()\n",
        "error_tree = hnfg.errorTree(cont_values)\n",
        "print(\"\\nError Tree (Nonlinear errors across modes):\")\n",
        "error_tree.print()\n",
        "\n",
        "# --- discretePosterior(Values) ---\n",
        "# Calculates P(M | X=x_cont), normalizing exp(-errorTree)\n",
        "# Includes contributions from purely discrete factors as well.\n",
        "discrete_posterior = hnfg.discretePosterior(cont_values)\n",
        "print(\"\\nDiscrete Posterior P(M | X=x):\")\n",
        "discrete_posterior.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Linearization\n",
        "\n",
        "The primary operation is `linearize`, which converts the graph to a `HybridGaussianFactorGraph` at a given linearization point (continuous `Values`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Linearizing at:\n",
            "Values with 2 values:\n",
            "Value x0: (gtsam::Pose2)\n",
            "(0, 0, 0)\n",
            "\n",
            "Value x1: (gtsam::Pose2)\n",
            "(1, 0, 0)\n",
            "\n",
            "\n",
            "Resulting HybridGaussianFactorGraph:\n",
            "\n",
            "size: 3\n",
            "Factor 0\n",
            "GaussianFactor:\n",
            "\n",
            "  A[x0] = [\n",
            "\t10, 0, 0;\n",
            "\t0, 10, 0;\n",
            "\t0, 0, 20\n",
            "]\n",
            "  b = [ 0 0 0 ]\n",
            "  No noise model\n",
            "\n",
            "Factor 1\n",
            "DiscreteFactor:\n",
            " f[ (d0,2), ]\n",
            " Choice(d0) \n",
            " 0 Leaf  0.8\n",
            " 1 Leaf  0.2\n",
            "\n",
            "Factor 2\n",
            "HybridGaussianFactor:\n",
            "Hybrid [x0 x1; d0]{\n",
            " Choice(d0) \n",
            " 0 Leaf :\n",
            "  A[x0] = [\n",
            "\t-10, -0, -0;\n",
            "\t0, -10, -10;\n",
            "\t-0, -0, -57.2957795\n",
            "]\n",
            "  A[x1] = [\n",
            "\t10, 0, 0;\n",
            "\t0, 10, 0;\n",
            "\t0, 0, 57.2957795\n",
            "]\n",
            "  b = [ -0 -0 -0 ]\n",
            "  No noise model\n",
            "scalar: -5.89658155\n",
            "\n",
            " 1 Leaf :\n",
            "  A[x0] = [\n",
            "\t-2, -0, -0;\n",
            "\t0, -2, -2;\n",
            "\t-0, -0, -5.72957795\n",
            "]\n",
            "  A[x1] = [\n",
            "\t2, 0, 0;\n",
            "\t0, 2, 0;\n",
            "\t0, 0, 5.72957795\n",
            "]\n",
            "  b = [ -0 -0 -0 ]\n",
            "  No noise model\n",
            "scalar: -0.375120634\n",
            "\n",
            "}\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Linearization point (often the current estimate)\n",
        "linearization_point = Values()\n",
        "linearization_point.insert(X(0), Pose2(0, 0, 0))\n",
        "linearization_point.insert(X(1), Pose2(1, 0, 0))\n",
        "\n",
        "print(\"\\nLinearizing at:\")\n",
        "linearization_point.print()\n",
        "\n",
        "# Perform linearization\n",
        "hgfg = hnfg.linearize(linearization_point)\n",
        "\n",
        "print(\"\\nResulting HybridGaussianFactorGraph:\")\n",
        "hgfg.print()\n",
        "# Note: NonlinearFactors become JacobianFactors.\n",
        "#       HybridNonlinearFactors become HybridGaussianFactors.\n",
        "#       DiscreteFactors remain DiscreteFactors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Restriction (`restrict`)\n",
        "\n",
        "Applies `restrict` to all factors in the graph for a given discrete assignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Restricted HybridNonlinearFactorGraph (D0=0):\n",
            "HybridNonlinearFactorGraph\n",
            " \n",
            "size: 2\n",
            "factor 0: PriorFactor on x0\n",
            "  prior mean:  (0, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.05];\n",
            "\n",
            "factor 1: BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.0174532925];\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Restrict the graph to the 'Grippy' mode (D0=0)\n",
        "assignment = DiscreteValues()\n",
        "assignment[D(0)] = 0\n",
        "restricted_hnfg = hnfg.restrict(assignment)\n",
        "\n",
        "print(\"\\nRestricted HybridNonlinearFactorGraph (D0=0):\")\n",
        "restricted_hnfg.print()\n",
        "# Note: The HybridNonlinearFactor is replaced by its D0=0 component (odom0).\n",
        "# The DiscreteFactor is removed as its variable is assigned."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
